import json
import gzip
import os
import shutil
import subprocess
from typing import *
from c2rust.roles import baseline_team, rag_team, naive
from c2rust.roles.baseline import CompilerMessageAnalyser
from metagpt.team import Team

RESULT_DIR = "./results/"

def stream_jsonl(filename: str) -> Iterable[Dict]:
    """
    Parses each jsonl line and yields it as a dictionary
    """
    if filename.endswith(".gz"):
        with open(filename, "rb") as gzfp:
            with gzip.open(gzfp, "rt") as fp:
                for line in fp:
                    if any(not x.isspace() for x in line):
                        yield json.loads(line)
    else:
        with open(filename, "r") as fp:
            for line in fp:
                if any(not x.isspace() for x in line):
                    yield json.loads(line)  

def hire_team(use_naive = False, use_rag=True, max_iters=5):
    def add_role(role_cls):
        if role_cls == CompilerMessageAnalyser:
            return CompilerMessageAnalyser(epoch_limit=max_iters)
        else:
            return role_cls()
    
    team_roles = []
    team_blueprint = None
    if use_naive:
        team_blueprint = naive
    elif use_rag:
        team_blueprint = rag_team
    else:
        team_blueprint = baseline_team
    
    for role_blueprint in team_blueprint:
        team_roles.append(add_role(role_blueprint))
    return team_roles              
                    
async def run_model(
    c_code,
    use_naive = False,
    use_rag = True,
    investment = 2.0,
    max_round = 10
):
    investment = float(investment)
    n_round = int(max_round)
    
    # remove `workspace` dir, if exists
    if os.path.exists("workspace"):
        shutil.rmtree("workspace")

    team = Team()
    team.hire(hire_team(use_naive=use_naive, use_rag=use_rag))
    team.invest(investment=investment)
    
    await team.run(n_round=n_round, idea=c_code)
    
    result = subprocess.run(["rustc", "main.rs"], capture_output=True, text=True)
    code_result = result.stderr
    # 返回 main.rs 中的内容
    with open('main.rs', 'r') as f:
        return (f.read(), code_result)